#include "TargetDisk.h"
#include <iostream>
#include <Windows.h>
#include <string>
#include <math.h>

using namespace std;

CTargetDisk::CTargetDisk(char* pszDiskName, EDeviceType eDevType) :	m_phDiskFile(NULL), 
																	m_pwLogicalDiskName(NULL), 
																	m_psTestsList(NULL), 
																	m_nTestsNum(0),
																	m_pwFileName(NULL),
																	m_eDevType(eDevType),
																	m_pNextTestSegment(NULL),
																	m_hCompletionQueue(NULL)
{
	m_pwLogicalDiskName = ConvertChar2Wstr(pszDiskName);
	m_psStopWatch = new CStopWatch();
}

CTargetDisk::~CTargetDisk(void)
{
	if (m_pwLogicalDiskName != NULL)
	{
		free(m_pwLogicalDiskName);
	}


	// TODO delete the whole list
}

bool CTargetDisk::PreTest(char  *pszFileName, EIOType eIoType)
{
	m_pwFileName = ConvertChar2Wstr(pszFileName);

	// Open the disk
	SetErrorMode(SEM_FAILCRITICALERRORS);

	if (eIoType == eWrite)
	{
		m_phDiskFile = CreateFile(m_pwFileName, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ |
			       FILE_SHARE_WRITE, NULL, OPEN_ALWAYS,
			       FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED, NULL);
	}
	else
	{
		// Reading file just if it exists
		m_phDiskFile = CreateFile(m_pwFileName, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ |
		       FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
		       FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED, NULL);
	}

	if (m_phDiskFile == INVALID_HANDLE_VALUE)
	{
		cout << "failed in creating file - error: " << GetLastError()<< endl;
		return false;
	}
	else
	{
		if (m_eDevType == ePhysical)
		{
			// Lock the disk and unmount it for physical writing
			
			// Lock the device

			DWORD  lpBytesReturned = 0;
			BOOL bRet = DeviceIoControl(m_phDiskFile,  
						FSCTL_LOCK_VOLUME,
						NULL, 0, // no input buffer
						NULL, 0, // no input buffer
						&lpBytesReturned,// # bytes returned
						NULL);  // asynchronous I/


			if(!bRet)
			{
				cout << "*** Error " << GetLastError()<< " closing disk" << endl;
				CloseHandle(m_phDiskFile);
				return false;
			}

//			cout << "Locked the device" << endl;

			// Dismount the disk
			bRet = DeviceIoControl(
				(HANDLE) m_phDiskFile,            // handle to a volume
				(DWORD) FSCTL_DISMOUNT_VOLUME,   // dwIoControlCode
				NULL,                        // lpInBuffer
				0,                           // nInBufferSize
				NULL,                        // lpOutBuffer
				0,                           // nOutBufferSize
				&lpBytesReturned,   // number of bytes returned
				NULL
				);

			if(!bRet)
			{
				cout << "*** Error " << GetLastError()<< " closing disk" << endl;
				CloseHandle(m_phDiskFile);
				return false;
			}

//			cout << "dismounted the device" << endl;

			// Get rid of dos device

			TCHAR szDriveLetter[3];
			szDriveLetter[1] = TEXT('E');
			szDriveLetter[1] = TEXT(':');
			szDriveLetter[2] = TEXT('\0');

			bRet = DefineDosDevice (DDD_RAW_TARGET_PATH|DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE,
				szDriveLetter, NULL);

			if(!bRet)
			{
				cout << "*** Error " << GetLastError()<< " get rid of dos device" << endl;
				CloseHandle(m_phDiskFile);
				return false;
			}

		}

		if (CreateIoCompletionPort(m_phDiskFile, m_hCompletionQueue, 0, 1)== NULL)
		{
			
		}
		else
		{
				

			// Prepare every test segment before the operation so it won't be included
			// in the time measurment
			CTestSegment*	psTempSegment = m_psTestsList;
			while (m_nTestsNum > 0 && psTempSegment != NULL)
			{
				if (psTempSegment)
				{
					// Prepare the test
					psTempSegment->PrepareTest(m_phDiskFile);
					psTempSegment = psTempSegment->m_psNextSegment;

				}
			}
		}
	}
	return true;
}

wchar_t* CTargetDisk::ConvertChar2Wstr(char  *pszSrc)
{
	int len = strlen(pszSrc)+1;
	wchar_t* wDest = new wchar_t[len];
	if ( wDest == 0 )
	return NULL;
	memset(wDest,0,len);
	::MultiByteToWideChar(  CP_ACP, NULL,pszSrc, -1, wDest,len );
	return wDest;
}

BOOL CTargetDisk::GetLogicalDiskPhysicalInfo()
{
	BOOL bResult = FALSE;

	bResult = GetDiskFreeSpace (m_pwLogicalDiskName, 
                                 &m_dwSectPerClust,
                                 &m_dwBytesPerSect, 
                                 &m_dwFreeClusters,
                                 &m_dwTotalClusters);

	if (!bResult)
	{
		cout <<"error: "<<  GetLastError(); 
		std::wcout << "could not get free space for "<< m_pwLogicalDiskName << endl;
	}
	else
	{
		std::wcout << "Device name: \t\t\t" << m_pwLogicalDiskName << endl;
		cout <<"Bytes per sectotr: \t\t" << m_dwBytesPerSect << endl;
		cout <<"Sectors per clusters: \t\t" << m_dwSectPerClust << endl;
		cout << "Total number of clusters: \t" <<  m_dwTotalClusters << endl;
		cout << "Number of free clusters : \t" << m_dwFreeClusters << endl;
		
		cout << endl;
		double fTotalSize = ((m_dwBytesPerSect * m_dwSectPerClust * m_dwTotalClusters) / pow(10.0, 9.0));
		cout << "Calculate device size: \t\t "<< fTotalSize <<" GB " <<endl;
		fTotalSize = ((m_dwBytesPerSect * m_dwSectPerClust * m_dwFreeClusters) / pow(10.0, 9.0));
		cout << "Calculate device size (free): \t "<< fTotalSize <<" GB " <<endl;

	}

	return bResult;
}

void CTargetDisk::PostTest()
{
	
	// 1. delete the tests list

	CTestSegment*	psNextTempSegment = m_psTestsList->m_psNextSegment;
	CTestSegment*	psCurrTempSegment = m_psTestsList;

	while (m_nTestsNum > 0 || psCurrTempSegment != NULL)
	{
		delete psCurrTempSegment;
		m_nTestsNum--;
		if (psNextTempSegment != NULL)
		{
			psCurrTempSegment = psNextTempSegment;
			psNextTempSegment = psCurrTempSegment->m_psNextSegment;
		}
		else
		{
			psCurrTempSegment = NULL;
		}
	}

	m_psTestsList = NULL;


	// 2. free the file name 
	if (m_pwFileName)
	{
		free(m_pwFileName);
	}

	if (m_eDevType == ePhysical)
	{
		DWORD dwBytesReturned = 0;
		// UnLock the disk
		BOOL bRet = DeviceIoControl(m_phDiskFile,  
			FSCTL_UNLOCK_VOLUME,
			NULL, 0, // no input buffer
			NULL, 0, // no input buffer
			&dwBytesReturned,// # bytes returned
			NULL);  // synchronous I/

		if (bRet != TRUE)
		{
			cout << "*** Error " << GetLastError()<< " unlocking disk" << endl;
		}
	}

	// 3. close the disk
	if (!CloseHandle(m_phDiskFile))
	{
		cout << "*** Error " << GetLastError()<< " closing disk" << endl;
		SetLastError(0);
	}
	else
	{
//		cout << "Succeeded in closing the disk" << endl;
	}

}

bool CTargetDisk::DoOverlappedIO(CTestSegment* psTest)
{
	DWORD dwBytesWritten = 0;
	BOOL bRes = FALSE;
	string sOp;
	EIOType eIoType = psTest->m_eIoType; 
	if (eIoType == eWrite)
	{
		bRes = WriteFile(m_phDiskFile, (char *)psTest->m_pBuffer, 
			psTest->m_dwSegmentBytes, &dwBytesWritten, &(psTest->m_oOverlapInfo));
//		sOp = "writing";
	}
	else
	{
		bRes = ReadFile(m_phDiskFile, (char *)psTest->m_pBuffer, 
			psTest->m_dwSegmentBytes, &dwBytesWritten, &(psTest->m_oOverlapInfo));
//		sOp = "reading";
	}

	// Do the asynchronous IO.
	if (bRes == TRUE) {
			// It succeeded immediately!
			//cout << sOp << "(immediately) " << dwBytesWritten << " of "
			//    << psTest->m_dwSegmentBytes << " bytes to disk ";
			//std::wcout << m_pwFileName << "." << endl;
	}
	else {
			if (GetLastError() != ERROR_IO_PENDING) 
			{ 
				 // WriteFile failed, but isn't delayed. Report error and abort.
				// It didn't start OK!  Bail out.
				cout << "***Error (immediately) " << sOp << dwBytesWritten <<
					" of " << psTest->m_dwSegmentBytes << " bytes to disk: ";
				std::wcout << m_pwFileName;
				cout <<", error=" << GetLastError() << endl;
				CloseHandle(psTest->m_oOverlapInfo.hEvent);
				return false;
			}
	}
	 // Write is pending.
     DWORD dwRes = WaitForSingleObject(psTest->m_oOverlapInfo.hEvent, INFINITE);
     switch(dwRes)
     {
        // OVERLAPPED structure's event has been signaled. 
        case WAIT_OBJECT_0:
		{

			// Check to see if it has completed
			if (GetOverlappedResult(m_phDiskFile, &(psTest->m_oOverlapInfo), &dwBytesWritten, FALSE)) 
			{
				// It completed successfully!
				//cout << sOp << " (eventually) " << dwBytesWritten << " of " << psTest->m_dwSegmentBytes
				//	<< " bytes to disk ";
				//std::wcout << m_pwFileName << "." << endl;
			}
			else
			{
				// It  wasn't completed successfully!
				cout << "***Error (eventually) " << sOp  << dwBytesWritten << " of " << psTest->m_dwSegmentBytes
					<< " bytes to disk ";
				std::wcout << m_pwFileName;
				cout << "."   << "error=" << GetLastError() << endl;
				CloseHandle(psTest->m_oOverlapInfo.hEvent);
				return false;
			}
		}

	 }
	CloseHandle(psTest->m_oOverlapInfo.hEvent);
	return true;

}

double CTargetDisk::PerformTest()
{
	bool bRes = false;
	CTestSegment*	psTempSegment = m_psTestsList;
	DWORD dwSegmentBytes = 0;

	DWORD dwStart = timeGetTime();
	m_psStopWatch->startTimer();

	int nTestsNum = m_nTestsNum;


	while (nTestsNum > 0 && psTempSegment != NULL)
	{
		if (psTempSegment)
		{
			// Perform the test
			bool bRes = DoOverlappedIO(psTempSegment);
			dwSegmentBytes = psTempSegment->m_dwSegmentBytes;
			psTempSegment = psTempSegment->m_psNextSegment;
			nTestsNum--;

		}
	}
	DWORD dwStop = timeGetTime();
	m_psStopWatch->stopTimer();
	DWORD dwTimeDiffMsc = dwStop - dwStart;
	double dElapsedTime = m_psStopWatch->getElapsedTime();

//	cout << "Time elapsed: " << dwTimeDiffMsc << " ms"<< endl;
	cout << "Time elapsed: " << dElapsedTime << " sec"<< endl;

	if (dwTimeDiffMsc > 0)
	{
		double fRate = ((dwSegmentBytes * 1000.0) / (1024.0*1024.0) / (dElapsedTime));

		cout << "Rates: " <<  fRate << "[MB/sec]" << endl;
	}

	return dElapsedTime;
//	return dwTimeDiffMsc;


}

void CTargetDisk::AddTest(CTestSegment* newSegmentTest)
{
	if (m_psTestsList == NULL)
	{
		m_psTestsList = newSegmentTest;
		m_pNextTestSegment = m_psTestsList;
	}
	else
	{
		m_pNextTestSegment->m_psNextSegment = newSegmentTest;
		m_pNextTestSegment = m_pNextTestSegment->m_psNextSegment;
	}
	m_nTestsNum++;
}

void CTargetDisk::CleanTestList()
{
	int nCount = m_nTestsNum;
	CTestSegment* psTemp = NULL;
	for (int i = 0; i < m_nTestsNum; ++i)
	{
		psTemp = m_psTestsList->m_psNextSegment;
		delete m_psTestsList;
		m_psTestsList = psTemp;
		m_nTestsNum--;
	}
	m_psTestsList = NULL;
}
